from requests import post
from time import time
from rich import print
from typing import (
    TypedDict,
    Literal
)

base_url = 'http://localhost:5000/predict'
text = '《长津湖》由陈凯歌、徐克、林超贤联合导演，吴京、易烊千玺主演。'

ModelName = Literal['BiLSTM', 'BiLSTM_CRF']

class InputText(TypedDict):
    text: str

def lstm_predict(input_text: InputText) -> None:
    time_list = []
    inputs = {'input_text': input_text}
    print('[bold gold1]Bi-LSTM & CRF[/]')
    for _ in range(10):
        start = time()
        resp = post(url=base_url, json=inputs)
        end = time()
        time_taken = (end - start) * 1000
        time_list.append(time_taken)
    avg_time = sum(time_list) / len(time_list)
    print(f'[bold green]time taken: [bold white]{avg_time:.2f}ms[/]')
    print(f'[bold green]result:[/] {resp.json()['result']}')

if __name__ == '__main__':
    input_text = {'text': text}
    lstm_predict(input_text)
